import pandas as pd
import numpy as np

# 初始化随机种子
np.random.seed(42)

# 生成基本参数
n = 116
data = pd.DataFrame(index=range(1, n+1))

# 1. 常居地（62%城市，38%农村）
data['常居地'] = np.random.choice(['城市', '农村'], n, p=[0.62, 0.38])

# 2. 月收入（按城乡分层）
def assign_income(residence):
    if residence == '农村':
        return np.random.choice(['<3000', '3000-5000', '5000-8000'], p=[0.12, 0.55, 0.33])
    else:
        return np.random.choice(['5000-8000', '8000-12000', '>12000'], p=[0.38, 0.45, 0.17])
data['月收入（元）'] = data['常居地'].apply(assign_income)

# 3. 教育水平（全局分布）
data['教育水平'] = np.random.choice(
    ['初中及以下', '高中/中专', '大学本科', '硕士及以上'], 
    n, 
    p=[0.18, 0.32, 0.42, 0.08]
)

# 4. 年龄（均值29.6岁，标准差4.0）
data['年龄'] = np.clip(np.round(np.random.normal(29.6, 4, n)), 19, 42).astype(int)

# 5. 初次妊娠（62%是）
data['初次妊娠'] = np.random.choice(['是', '否'], n, p=[0.62, 0.38])

# 6. 不良孕史（非初次妊娠组76%有，初次组8.3%有）
def assign_pregnancy_history(row):
    if row['初次妊娠'] == '否':
        return np.random.choice(['是', '否'], p=[0.76, 0.24])
    else:
        return np.random.choice(['是', '否'], p=[0.083, 0.917])
data['不良孕史'] = data.apply(assign_pregnancy_history, axis=1)

# 7. 睡眠质量（城乡不同均值）
def assign_sleep(residence):
    if residence == '农村':
        return np.clip(np.round(np.random.normal(10.2, 4.1)), 0, 21)
    else:
        return np.clip(np.round(np.random.normal(13.6, 4.1)), 0, 21)
data['睡眠质量满意度'] = data['常居地'].apply(assign_sleep)

# 8. 易感性人格量表（均30.7，标准差7.8）
data['易感性人格量表'] = np.clip(np.round(np.random.normal(30.7, 7.8, n)), 9, 45).astype(int)

# 9. EPDS评分（与睡眠、易感性相关，并发症组更高）
data['妊娠并发症'] = np.random.choice(['无', '有'], n, p=[0.81, 0.19])
data['EPDS评分'] = np.clip(
    (0.6 * data['易感性人格量表'] - 0.5 * data['睡眠质量满意度'] + np.random.normal(0, 3, n)),
    0, 30
).astype(int)
data.loc[data['妊娠并发症'] == '有', 'EPDS评分'] += np.random.randint(3, 7, sum(data['妊娠并发症'] == '有'))

# 10. 参访状态（84.5%参访）
data['参访状态'] = np.random.choice(['参访', '失访'], n, p=[0.845, 0.155])

# 输出
print(data.head(15))
data.to_csv("pregnancy_survey_data.csv", index_label="序号",encoding='gbk')